#from libmyatomism import *
#from math import *
#export LD_PRELOAD=/usr/lib/libmpi.so python
import array 
import matplotlib.pyplot as plt
from libatomism import *
import math
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.mlab import griddata
from matplotlib import cm


""" this function plot the 2D PES isocontours, for your 
first lecture you should go directly to the main body
of the code.
"""
def plotIsoPesForApproach( kinop, repro, fig , nfig):

   zmat = kinop.getSystem()
   gcoors = kinop.getCoordinates()
   q0 = gcoors[0]
   q1 = gcoors[1]
   q0.set(2.,2.,5.,0.1,0.1)
   q1.set(0.1,0.1,90.1,1,1)

   """ A linear sampler: scan from min to max like a 'for' loop
   imbriqued for each activated coordinates (if Dq>0)
   """
   scan=msSamplerLinear.New()
  
   gcoors[2].Dq=0
   gcoors[3].Dq=0
   """ the iso levels we want to plot """
   levels=[-0.9,-0.75,-0.5,-0.25,0,2.5,5,10,20,40,50,60,70,80]
   zeroindice=4
   
   """ create the array which will contain the values of the 2D PES
   """
   X = np.arange(q0.Min,q0.Max,q0.Dq)
   Y = np.arange(q1.Min,q1.Max,q1.Dq)
   Z = np.ones( (len(Y),len(X) ))  
   X, Y = np.meshgrid(X, Y)

   scan.begin()
   while( scan.getNextPoint( gcoors ) ):
  
       """ set the coordinates and the atoms' position
       """    
       kinop.setDynamicDofs()
       zmat.computeCartCoordinates()    

       """ Here we get the values of the qij_repro
       """
       vec = VectorOfDouble()
       vec.extend([zmat.separation(1,4)**(-2),zmat.separation(1,5)**(-2),zmat.separation(2,4)**(-2),zmat.separation(2,5)**(-2)])
       a =  repro.evaluate(vec)

       if( (a>Emax) or (a<-10) ):
           a=Emax
       Z[q1.indice()][q0.indice()] = a
   """
   plot
   """
   plt=fig.add_subplot(2, 2, nfig)
   CS = plt.contour(X, Y, Z,levels, colors='k')
   zc = CS.collections[zeroindice] 
   zc.set_linewidth(4)
   plt.clabel(CS, inline=1, fontsize=12)
   plt.set_xlabel("q0 [Angstrom]",fontsize=20)
   plt.set_ylabel("q1 [Degree]",fontsize=20)
   plt.tick_params(labelsize=20)


""" main body of the code
"""
msLogger.setPriorityMin(Priority.ERROR)

unit=msUnitsManager.New("Angstrom Degree amu kcal/mol ps")

"""
Declaration of the z-matrix
X 
N 0 0.55 
N 0 0.55        1  179.9  
X 0 5.   (=q0)  1  90.   (=q1)   2   0. 
N 3 0.55        0  90.   (=q2)   1   180. (=q3)
N 3 0.55        4  179.9         0   180.
"""
zmat = msZMat.New(unit).setId("zmat").set("X N 0 0.55 N 0 0.55 1 179.9 X 0 5. 1 90. 2 0. N 3 0.55 0 90. 1 180. N 3 0.55 4 179.9 0 180.")

system=msSystem.New(unit).addEntity(zmat)

"""
Declaration of the dynamic variables of the z-matrix, 
those who will control the relative orientation of the two N2.
"""
q0 = msScalarVariable.New("Angstrom").set(5,0.5,6,0.1,0.1).setId("q0")
q1 = msScalarVariable.New("Degree").set(120,0.1,179.9,2,2).setId("q1")
q2 = msScalarVariable.New("Degree").set(120,0.1,179.9,2,2).setId("q2")
q3 = msScalarVariable.New("Degree").set(120,0,360,2,2).setId("q3")

gcoors = msGeneralizedCoordinates.New(unit).addVar(q0).addVar(q1).addVar(q2).addVar(q3)

"""
The kinetic operator: it does the link between generalized 
coordinates and geometry
"""
kinop=msKineticOperator.New(unit).set(gcoors,system)

"""
the following is defining the kinetic operator with dynamic DoF of type f(q)=q;
the first definition is the same as doing, for instance:
  kinfct0=msParser.New("Angstrom").setCoordinates(gcoors).setExpression("q0").setId("f=q0")
  kinop.addDynamicDof(zmat,3,kinfct0)
"""
kinop.addDynamicDofLinear(zmat,3,q0).addDynamicDofLinear(zmat,4,q1).addDynamicDofLinear(zmat,7,q2).addDynamicDofLinear(zmat,8,q3)

""" 
Load the database ( format of each line:  q0 q1 q2 q3 q4 E_in_hartree E_in_kcal/mol ) 
we want to fit E_in_kcal/mol (0 of energy is 2*E(N2) )
"""
with open('Database.txt') as f:
    w = f.readline().split() # read first line
    array = []
    for line in f:           # read rest of lines
        array.append([float(x) for x in line.split()])

""" 
Define the variable of the fitting function: 
this is not the gcoors because we want to use variable of type 1/(R*R)
where R is the atoms separation. For instance, q02 refer to the separation
 between the nitrogen atom number O and number 2.
"""
q02_repro = msScalarVariable.New("Angstrom^-2").set(0,0,1.e6,0.1,0.1).setId("q0_repro")
q03_repro = msScalarVariable.New("Angstrom^-2").set(0,0,1.e6,0.1,0.1).setId("q1_repro")
q12_repro = msScalarVariable.New("Angstrom^-2").set(0,0,1.e6,0.1,0.1).setId("q2_repro")
q13_repro = msScalarVariable.New("Angstrom^-2").set(0,0,1.e6,0.1,0.1).setId("q3_repro")
gcoors_repro = msGeneralizedCoordinates.New(unit).addVar(q02_repro).addVar(q03_repro).addVar(q12_repro).addVar(q13_repro)

"""
Repromodel is a type of scalar function, we initialize it 
with the unit of the returned value, and set the coordinates after that.
""" 
repro =msRepromodel.New("kcal/mol").setCoordinates(gcoors_repro)

"""
This is the order of the development, and the maximum energy accounted
""" 
N = 5
Emax=80

"""
Here we add the data to fit into the repromodel.
The database inside the repromodel is a scalar function of 
type dirac.
"""
for values in array:

    vec = VectorOfDouble()
    vec.extend( arg for arg in values)
    """ set the coordinates and the atoms' position
    """
    gcoors.setCoordinates(vec)
    kinop.setDynamicDofs()
    system.computeCartCoordinates()
    """ Here we get the values of the qij_repro
    """
    vec = VectorOfDouble()
    vec.extend([ zmat.separation(1,4)**(-2) , zmat.separation(1,5)**(-2) , zmat.separation(2,4)**(-2) , zmat.separation(2,5)**(-2)])
    if( values[5]<Emax ):
        repro.addData( vec, values[5] )
        
msLogger.setPriorityMin(Priority.INFO)
msLogger.setVerboseLevel(1)

""" Create the orthonormal basis set of order N
"""
repro.setBasisSet(N)
msLogger.setPriorityMin(Priority.ERROR)

""" fit the function: projection of the database function (i.e dirac function)
on the basis set. 
"""
repro.fit()

""" After this point, this is only for plotting the results ...

return a sampler containing the scenario 
of the database providen
"""
sampler = repro.getSampler()

""" get the database: it is a scalar
function of type dirac, each time evaluated 
for the sampler current position. see the 
tutorial page for details.
"""
field = repro.getField()

y=[]
x=[]
n=0
sigma=0

""" begin the sampling """
sampler.begin()

"""comparison between extact and fitted values
"""
while(sampler.getNextPoint(gcoors_repro)):
     x.append(field.evaluate())
     y.append(repro.evaluate())
     sigma+=abs(field.evaluate() - repro.evaluate())
     n=n+1
print "sigma=",sigma/n


fig = plt.figure()
plt=fig.add_subplot(2, 2, 1)
plt.axis((-5,Emax,-5,Emax))
plt.set_xlabel("Exact energy [kcal/mol]",fontsize=20)
plt.set_ylabel("Fitted energy [kcal/mol]",fontsize=20)
plt.tick_params(labelsize=20)

plt.plot(x,y,"+")


""" First approach
"""
q2.Value=179.9
q3.Value=0.
plotIsoPesForApproach( kinop, repro, fig , 2)

""" same for the second approach
"""
q2.Value=90.
plotIsoPesForApproach( kinop, repro, fig , 3)

""" and the third approach
"""
q2.Value=90.
q3.Value=90.

plotIsoPesForApproach( kinop, repro, fig , 4)

fig.set_size_inches(20,20)
fig.savefig('N4_0.png')

fig.show()  
